
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.optim import Adam
from torch.nn import Linear
from torch.utils.data import DataLoader
from model.cluster.SDCN.GNN import GNNLayer
from model.cluster.SDCN.utils import construct_graph
from sklearn.cluster import KMeans
from torch.utils.tensorboard import SummaryWriter

torch.manual_seed(66)  # 控制随机状态
torch.cuda.manual_seed(66)  # 控制随机状态
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class AE(nn.Module):
    """
    自编码器类
    """

    def __init__(self, n_enc_1, n_enc_2, n_enc_3, n_dec_1, n_dec_2, n_dec_3,
                 n_input, n_z):
        """
        初始化自编码器结构
        Args:
            n_enc_1: 第一层编码的隐含层节点个数
            n_enc_2: 第二层编码的隐含层节点个数
            n_enc_3: 第三层编码的隐含层节点个数
            n_dec_1: 第一层解码的隐含层节点个数
            n_dec_2: 第二层解码的隐含层节点个数
            n_dec_3: 第三层解码的隐含层节点个数
            n_input: 输入数据的特征
            n_z: 编码之后的数据特征
        """
        super(AE, self).__init__()
        self.enc_1 = Linear(n_input, n_enc_1)
        self.enc_2 = Linear(n_enc_1, n_enc_2)
        self.enc_3 = Linear(n_enc_2, n_enc_3)
        self.z_layer = Linear(n_enc_3, n_z)

        self.dec_1 = Linear(n_z, n_dec_1)
        self.dec_2 = Linear(n_dec_1, n_dec_2)
        self.dec_3 = Linear(n_dec_2, n_dec_3)
        self.x_bar_layer = Linear(n_dec_3, n_input)

    def forward(self, x):
        """
        网络前向传递
        Returns:
            x_bar：数据重构
            enc_h1：第一层编码数据
            enc_h2：第二层编码数据
            enc_h3：第三层编码数据
            z：最后一层编码数据
        """
        enc_h1 = F.relu(self.enc_1(x))
        enc_h2 = F.relu(self.enc_2(enc_h1))
        enc_h3 = F.relu(self.enc_3(enc_h2))
        z = self.z_layer(enc_h3)

        dec_h1 = F.relu(self.dec_1(z))
        dec_h2 = F.relu(self.dec_2(dec_h1))
        dec_h3 = F.relu(self.dec_3(dec_h2))
        x_bar = self.x_bar_layer(dec_h3)

        return x_bar, enc_h1, enc_h2, enc_h3, z


class SDCN(nn.Module):
    """
    深度结构聚类网络
    """

    def __init__(self, n_enc_1, n_enc_2, n_enc_3, n_dec_1, n_dec_2, n_dec_3,
                 n_input, n_z, n_clusters, v=1.0):
        """
        初始化网络结构
        Args:
            n_enc_1: 第一层编码的隐含层节点个数
            n_enc_2: 第二层编码的隐含层节点个数
            n_enc_3: 第三层编码的隐含层节点个数
            n_dec_1: 第一层解码的隐含层节点个数
            n_dec_2: 第二层解码的隐含层节点个数
            n_dec_3: 第三层解码的隐含层节点个数
            n_input: 输入数据的特征
            n_z: 编码之后的数据特征
            n_clusters: 聚类簇数
            v: degree
        """
        super(SDCN, self).__init__()

        # autoencoder for intra information
        self.ae = AE(
            n_enc_1=n_enc_1,
            n_enc_2=n_enc_2,
            n_enc_3=n_enc_3,
            n_dec_1=n_dec_1,
            n_dec_2=n_dec_2,
            n_dec_3=n_dec_3,
            n_input=n_input,
            n_z=n_z)

        # GCN for inter information
        # self.gnn_1 = GNNLayer(n_input, n_z)
        self.gnn_1 = GNNLayer(n_input, n_enc_1)
        self.gnn_2 = GNNLayer(n_enc_1, n_enc_2)
        self.gnn_3 = GNNLayer(n_enc_2, n_enc_3)
        self.gnn_4 = GNNLayer(n_enc_3, n_z)
        self.gnn_5 = GNNLayer(n_z, n_clusters)

        # cluster layer
        self.cluster_layer = Parameter(torch.Tensor(n_clusters, n_z))
        torch.nn.init.xavier_normal_(self.cluster_layer.data)

        # degree
        self.v = v

    def forward(self, x, adj):
        """
        网络前向传递机制
        Args:
            x: 聚类数据，shape->（样本数，特征）
            adj: 数据对应图的的邻接矩阵
        Returns:
            x_bar: 重构数据
            q：Dual Self-supervised Module
            predict：预测结果，shape->（样本数，簇数）
            z：编码数据
        """
        # DNN Module
        x_bar, tra1, tra2, tra3, z = self.ae(x)

        sigma = 0.5

        # GCN Module
        h = self.gnn_1(x, adj)
        h = self.gnn_2((1 - sigma) * h + sigma * tra1, adj)
        h = self.gnn_3((1 - sigma) * h + sigma * tra2, adj)
        h = self.gnn_4((1 - sigma) * h + sigma * tra3, adj)
        h = self.gnn_5((1 - sigma) * h + sigma * z, adj, active=False)
        pre = F.softmax(h, dim=1)

        # Dual Self-supervised Module
        q = 1.0 / (1.0 + torch.sum(torch.pow(z.unsqueeze(1) -
                                             self.cluster_layer, 2), 2) / self.v)
        q = q.pow((self.v + 1.0) / 2.0)
        q = (q.t() / torch.sum(q, 1)).t()

        return x_bar, q, pre, z


def target_distribution(q):
    """
    Args:
        q: Dual Self-supervised Module 返回值
    Returns:
        p: 目标分布
    """
    weight = q**2 / q.sum(0)
    return (weight.t() / weight.sum(1)).t()


def adjust_learning_rate(optimizer, epoch):
    """
    调整学习率
    Args:
        optimizer: 优化器
        epoch: int，训练epoch
    """
    lr = 0.001 * (0.1 ** (epoch // 20))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def pretrain_ae(data, n_input, n_z,
                n_enc_1=888,
                n_enc_2=464,
                n_enc_3=128,
                n_dec_1=128,
                n_dec_2=464,
                n_dec_3=888,
                lr = 0.001,
                EPOCH = 50):
    """
    对AE模块进行预训练
    Args:
        data: 聚类数据
        n_input: 输入特征
        n_z: 编码特征
        n_enc_1: 第一层编码的隐含层节点个数
        n_enc_2: 第二层编码的隐含层节点个数
        n_enc_3: 第三层编码的隐含层节点个数
        n_dec_1: 第一层解码的隐含层节点个数
        n_dec_2: 第二层解码的隐含层节点个数
        n_dec_3: 第三层解码的隐含层节点个数
        lr: 学习率
        EPOCH: 训练轮数

    Returns:

    """
    model = AE(n_enc_1, n_enc_2, n_enc_3, n_dec_1, n_dec_2, n_dec_3,
               n_input, n_z).to(device)
    train_loader = DataLoader(data, batch_size=128, shuffle=True)
    print(model)
    optimizer = Adam(model.parameters(), lr=lr)
    for epoch in range(EPOCH):
        # adjust_learning_rate(optimizer, epoch)
        for batch, x in enumerate(train_loader):
            x = x.to(device)

            x_bar, _, _, _, _ = model(x)
            loss = F.mse_loss(x_bar, x)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        with torch.no_grad():
            x = torch.Tensor(data).to(device).float()
            x_bar, _, _, _, _ = model(x)
            loss = F.mse_loss(x_bar, x)
            print('预训练自编码器：{} loss: {}'.format(epoch, loss))
    torch.save(model.state_dict(), 'ae.pkl')


def train_sdcn(data, adj, n_input, n_z, n_clusters,
               n_enc_1=512,
               n_enc_2=128,
               n_enc_3=128,
               n_dec_1=128,
               n_dec_2=128,
               n_dec_3=512,
               v=1.0,
               lr=0.001,
               epoch=200,
               is_save_model=True,
               save_path=r'sdcn.pkl',
               is_writer=False
               ):
    """
    训练SDCN模型
    Args:
        data: np.array，data_type:float32, 聚类数据
        adj: np.array，data_type:float32, 数据对应的邻接矩阵
        n_input: 数据的特征数
        n_z: 编码之后的调整数
        n_clusters: 聚类簇数
        lr: 学习率
        epoch: 训练轮数
        is_save_model: 是否保存模型
        save_path: 训练完成之后模型的保存路径
        n_enc_1: 第一层编码的隐含层节点个数
        n_enc_2: 第二层编码的隐含层节点个数
        n_enc_3: 第三层编码的隐含层节点个数
        n_dec_1: 第一层解码的隐含层节点个数
        n_dec_2: 第二层解码的隐含层节点个数
        n_dec_3: 第三层解码的隐含层节点个数
        v: degree
        is_writer: 是否需要监测训练过程中模型的权重和梯度数据
    """
    model = SDCN(n_enc_1,
                 n_enc_2,
                 n_enc_3,
                 n_dec_1,
                 n_dec_2,
                 n_dec_3,
                 n_input=n_input,
                 n_z=n_z,
                 n_clusters=n_clusters,
                 v=v).to(device)
    try:
        model.ae.load_state_dict(torch.load('ae.pkl'))
    except:
        pretrain_ae(data, n_input, n_z,
                    n_enc_1,
                    n_enc_2,
                    n_enc_3,
                    n_dec_1,
                    n_dec_2,
                    n_dec_3,
                    lr=lr,
                    EPOCH=10,
                    )
        model.ae.load_state_dict(torch.load('ae.pkl'))
    print(model)

    if is_writer:
        writer = SummaryWriter()

    optimizer = Adam(model.parameters(), lr=lr)

    data = torch.from_numpy(data).to(device)
    adj = torch.from_numpy(adj).to(device)
    with torch.no_grad():
        _, _, _, _, z = model.ae(data)
    kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=66)
    kmeans.fit(z.data.cpu().numpy())
    model.cluster_layer.data = torch.tensor(kmeans.cluster_centers_).to(device)

    for e in range(epoch):
        x_bar, q, pred, _ = model(data, adj)
        p = target_distribution(q.data)
        kl_loss = F.kl_div(q.log(), p, reduction='batchmean')
        ce_loss = F.kl_div(pred.log(), p, reduction='batchmean')
        re_loss = F.mse_loss(x_bar, data)
        loss = 0.1 * kl_loss + 0.01 * ce_loss + re_loss
        print(f'epoch:{e} | loss:{loss.item()}')

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if is_writer:
            for name, param in model.named_parameters():
                writer.add_histogram(name+'_weight', param.data, global_step=e)
                writer.add_histogram(name+'_grad', param.grad, global_step=e)

    if is_writer:
        writer.close()

    if is_save_model:
        torch.save(model, save_path)


def predict(data, adj, model_path=r'sdcn.pkl'):
    """
    读取训练好的模型进行测试
    Args:
        data: np.array，data_type:float32, 聚类数据
        adj: np.array，data_type:float32, 数据对应的邻接矩阵
        model_path: 模型的保存路径
    Returns:
        predict: 预测的属于每个簇的概率值（假概率）
        index: 预测所属于的簇编号
    """
    with torch.no_grad():
        data = torch.from_numpy(data).to(device)
        adj = torch.from_numpy(adj).to(device)
        model = torch.load(model_path)
        _, _, pred, _ = model(data, adj)
        _, ind = torch.max(pred, dim=1)

    return pred.cpu().data.numpy(), ind.cpu().data.numpy()


if __name__ == "__main__":
    import numpy as np
    data = np.random.randint(0, 100, (100, 50)).astype(np.float32)
    adj = np.ones((100, 100)).astype(np.float32)
    knn_graph = construct_graph(data).astype(np.float32)
    train_sdcn(data, knn_graph, 50, 10, 10)
